import numpy as np
import json
in_sub = [1,4,5,6,56,57,59,61,62,63,153,156,158,160,164,166,167,170,172,198,209,243,250,331,332,334,375,376,377,378,379,381]
test_in_sub = [6,334,375,376,378,379]
sub = [0,2,3,58,60,67,70,71,72,75,77,80,84,85,88,90,92,94,95,96,99,105,106,107,108,109,110,111,112,113,114,115,116,117,118,119,\
       120,121,145,154,155,157,159,161,162,163,165,168,169,171,173,176,177,179,180,183,184,189,190,192,194,196,200,201,202,204,\
       206,207,212,246,247,249,251,253,254,256,259,261,262,264,267,268,271,273,274,276,277,280,282,283,284,287,290,292,293,295,\
       296,297,298,299,300,301,302,303,304,305,306,307,308,309,310,311,312,313,329,330,333,335,374,380,382,384,387,388,391]
test_sub = [388,391,71,301,310,96,284,298,307,308,313,75,105,274,305,374,251,99,0,304,72,177,80,249]
test_add = [34,52,368,291,66,230,231,385,215,270,347,42,181,315,82,144,69,223\
,240,103,218,281,15,255,324,390,27,186,336,213,185,41,238,23,219,360\
,365,263,345,208,257,394,29,123,31,13,139,314,182]
l = [170,332,334,375,377,389,381,335,380,384,388,391,317,319,322,325,328,318,320,323,327,336,338,340,343,347,349,352,355,337,339,341,346,348,350,354,356,357,358,359,362,363,365,366,369,6,71,301,96,284,298,307,308,313,75,99,0,177,90,274,305,109,162,60,23,182,13,209,257,370,371,373,383,385, 386, 390,393, 181, 123]
def write_fold(filename, l):
    with open(filename, 'w') as f:
        json.dump(l, f) 
t = []
for elem in range(395):
    if elem not in l:
        t.append(elem)
t = list(np.random.permutation(t))

total_l = t+l
seg = [0]
for i in range(5):
    seg.append(int(395*0.2*(i+1)))
print seg
print len(total_l)
test_l = []
train_l = []
for i in range(5):
   test_l = total_l[seg[i]:seg[i+1]] 
   train_l = total_l[0:seg[i]] + total_l[seg[i+1]:]
   write_fold("test_fold_"+str(i)+".json",test_l)     
   write_fold("train_fold_"+str(i)+".json",train_l)     
   



